package com.smallstep.apilimit.interceptor;

import com.google.common.util.concurrent.RateLimiter;
import com.smallstep.apilimit.properties.ApiLimitGroupProperties;
import com.smallstep.apilimit.properties.ApiLimitHeaderPatternProperties;
import com.smallstep.apilimit.properties.ApiLimitProperties;
import com.smallstep.apilimit.properties.ApiLimitUrlPatternProperties;
import com.smallstep.apilimit.ratelimiter.RedisRateLimiter;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.util.AntPathMatcher;
import org.springframework.util.CollectionUtils;
import org.springframework.web.servlet.handler.HandlerInterceptorAdapter;

import javax.annotation.PostConstruct;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.*;
import java.util.regex.Pattern;

/**
 * @author 李俊
 * @Description
 * @Date 2019/9/3 15:11
 */
public class ApiLimitInterceptor extends HandlerInterceptorAdapter {

    @Autowired
    StringRedisTemplate stringRedisTemplate;

    private List<ApiLimitInterceptor.ApiLimitInfo> apiLimitInfoList = new ArrayList();
    @Autowired
    private ApiLimitProperties apiLimitProperties;

    public ApiLimitInterceptor() {

    }

    @PostConstruct
    private void init() {
        if (apiLimitProperties.getEnabled()) {
            Map<String, ApiLimitGroupProperties> settings = apiLimitProperties.getSettings();
            for (Map.Entry<String,ApiLimitGroupProperties> entry : settings.entrySet()) {
                ApiLimitGroupProperties apiLimitGroup = entry.getValue();
                if (CollectionUtils.isEmpty(apiLimitGroup.getUrlPatterns())
                        && CollectionUtils.isEmpty(apiLimitGroup.getHeaderPatterns())) {
                    throw new RuntimeException("ApiLimit config error, urlPatterns and headerPatterns not allowed to be empty at the same time!");
                } else {
                    ApiLimitInterceptor.ApiLimitInfo apiLimitInfo =  new ApiLimitInfo();
                    apiLimitInfo.setConfigId(entry.getKey());
                    apiLimitInfo.setPriority(apiLimitGroup.getPriority());
                    if (apiLimitProperties.isDistribute()) {
                        RedisRateLimiter redisRateLimiter = RedisRateLimiter.create(entry.getKey(),
                                (double)apiLimitGroup.getLimit(),stringRedisTemplate);
                        apiLimitInfo.setRedisRateLimiter(redisRateLimiter);
                    } else {
                        RateLimiter rateLimiter = RateLimiter.create((double)apiLimitGroup.getLimit());
                        apiLimitInfo.setRateLimiter(rateLimiter);
                    }
                    apiLimitInfo.setUrlPatterns(apiLimitGroup.getUrlPatterns());
                    apiLimitInfo.setHeaderPatterns(apiLimitGroup.getHeaderPatterns());
                    apiLimitInfoList.add(apiLimitInfo);
                }
                Collections.sort(apiLimitInfoList);
            }
        }
    }


    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if (!apiLimitProperties.getEnabled()) {
            return true;
        } else {
            //获取URI
            String uri = request.getRequestURI();
            //获取头部MAP信息
            Map<String,String> headMap = new HashMap<>();
            Enumeration<String> headNames = request.getHeaderNames();
            while (headNames.hasMoreElements()) {
                String headName = headNames.nextElement();
                headMap.put(headName,request.getHeader(headName));
            }
            String method = request.getMethod().toUpperCase();
            for (int i = 0; i < apiLimitInfoList.size(); i++) {
                ApiLimitInfo apiLimitInfo = apiLimitInfoList.get(i);
                this.checkLimit(apiLimitInfo,method,uri,headMap);
            }

        }
        return true;
    }

    /**
     *
     * @param apiLimitInfo
     * @param uri
     * @param headMap
     */
    private void checkLimit(ApiLimitInfo apiLimitInfo, String method, String uri, Map<String, String> headMap) {
        if (isUriMatch(apiLimitInfo,method,uri) && isHeaderMatch(apiLimitInfo,headMap)) {
            if (apiLimitProperties.isDistribute()) {
                if (!apiLimitInfo.getRedisRateLimiter().tryAcquire()) {
                    throw new RuntimeException("Interface access is forbidden to reach the maximum page view！");
                }
            } else {
                if (!apiLimitInfo.getRateLimiter().tryAcquire()) {
                    throw new RuntimeException("Interface access is forbidden to reach the maximum page view！");
                }
            }

        }

    }

    /**
     * 判断头部是否匹配
     * @param apiLimitInfo
     * @param headMap
     * @return
     */
    private boolean isHeaderMatch(ApiLimitInfo apiLimitInfo, Map<String, String> headMap) {
        List<ApiLimitHeaderPatternProperties> headerPatterns = apiLimitInfo.getHeaderPatterns();
        if (CollectionUtils.isEmpty(headerPatterns)) {
            return true;
        }
        for (int i = 0; i < headerPatterns.size(); i++) {
            ApiLimitHeaderPatternProperties headerPattern = headerPatterns.get(i);
            if (!headerPattern.isExists() && !headMap.containsKey(headerPattern.getKey())) {
                continue;
            }
            if (headerPattern.isExists() && headMap.containsKey(headerPattern.getKey())
            && Pattern.matches(headerPattern.getValuePattern(),headMap.get(headerPattern.getKey()))) {
                continue;
            }
            return false;
        }
        return true;
    }

    /**
     * 判断URI是否匹配
     * @param apiLimitInfo
     * @param uri
     * @return
     */
    private boolean isUriMatch(ApiLimitInfo apiLimitInfo, String method, String uri) {
        List<ApiLimitUrlPatternProperties> urlPatterns = apiLimitInfo.getUrlPatterns();
        if (CollectionUtils.isEmpty(urlPatterns)) {
            return true;
        }
        AntPathMatcher antPathMatcher = new AntPathMatcher();
        for (int i = 0; i < urlPatterns.size(); i++) {
            ApiLimitUrlPatternProperties urlPattern = urlPatterns.get(i);
            if (method.equals(urlPattern.getMethod()) &&
                    antPathMatcher.match(urlPattern.getPathPattern(),uri)) {
                return true;
            }
        }
        return false;
    }


    static class ApiLimitInfo implements Comparable<ApiLimitInfo>{
        /**
         * 优先级
         */
        private int priority;
        /**
         * 配置ID
         */
        private String configId;
        /**
         * 限流器
         */
        private RateLimiter rateLimiter;
        /**
         *
         */
        private RedisRateLimiter redisRateLimiter;
        /**
         * 路径匹配
         */
        private List<ApiLimitUrlPatternProperties> urlPatterns;
        /**
         * 头部匹配
         */
        private List<ApiLimitHeaderPatternProperties> headerPatterns;

        public int getPriority() {
            return priority;
        }

        public void setPriority(int priority) {
            this.priority = priority;
        }

        public String getConfigId() {
            return configId;
        }

        public void setConfigId(String configId) {
            this.configId = configId;
        }

        public RateLimiter getRateLimiter() {
            return rateLimiter;
        }

        public void setRateLimiter(RateLimiter rateLimiter) {
            this.rateLimiter = rateLimiter;
        }

        public List<ApiLimitUrlPatternProperties> getUrlPatterns() {
            return urlPatterns;
        }

        public void setUrlPatterns(List<ApiLimitUrlPatternProperties> urlPatterns) {
            this.urlPatterns = urlPatterns;
        }

        public List<ApiLimitHeaderPatternProperties> getHeaderPatterns() {
            return headerPatterns;
        }

        public void setHeaderPatterns(List<ApiLimitHeaderPatternProperties> headerPatterns) {
            this.headerPatterns = headerPatterns;
        }

        public RedisRateLimiter getRedisRateLimiter() {
            return redisRateLimiter;
        }

        public void setRedisRateLimiter(RedisRateLimiter redisRateLimiter) {
            this.redisRateLimiter = redisRateLimiter;
        }

        /**
         * 比较大小
         */
        @Override
        public int compareTo(ApiLimitInfo o) {
            //倒序
            return o.getPriority() - this.getPriority();
        }
    }
}
